from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import declarative_base
from contextlib import asynccontextmanager
from typing import AsyncGenerator
from config.db_config import get_config
from loguru import logger

Base = declarative_base()

class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self):
        self.engine = None
        self.session_factory = None
        self._initialized = False
    
    async def init_db(self):
        """初始化数据库连接"""
        if self._initialized:
            return
            
        try:
            config = get_config()
            db_config = config['database']
            
            # 构建SQLite数据库URL
            database_url = f"sqlite+aiosqlite:///{db_config['path']}"
            
            # 创建异步引擎
            self.engine = create_async_engine(
                database_url,
                echo=False  # 生产环境设为False
            )
            
            # 创建会话工厂
            self.session_factory = async_sessionmaker(
                bind=self.engine,
                class_=AsyncSession,
                expire_on_commit=False
            )
            
            self._initialized = True
            logger.info("数据库连接初始化成功")
            
        except Exception as e:
            logger.error(f"数据库连接初始化失败: {str(e)}")
            raise
    
    async def close_db(self):
        """关闭数据库连接"""
        if self.engine:
            await self.engine.dispose()
            self._initialized = False
            logger.info("数据库连接已关闭")
    
    @asynccontextmanager
    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取数据库会话"""
        if not self._initialized:
            await self.init_db()
            
        async with self.session_factory() as session:
            try:
                yield session
            except Exception:
                await session.rollback()
                raise
            finally:
                await session.close()

# 全局数据库管理器实例
db_manager = DatabaseManager()

# 便捷函数
@asynccontextmanager
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
    """获取数据库会话"""
    async with db_manager.get_session() as session:
        yield session

# 初始化和清理函数
async def init_database():
    """初始化数据库"""
    await db_manager.init_db()

async def close_database():
    """关闭数据库"""
    await db_manager.close_db()